{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from einops import rearrange, repeat\n",
    "from einops.layers.torch import Rearrange"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ResNet\n",
    "\n",
    "https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/Classification/ConvNets/image_classification/models/resnet.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SqueezeAndExcitation(nn.Module):\n",
    "    def __init__(self, in_channels, squeeze):\n",
    "        super(SqueezeAndExcitation, self).__init__()\n",
    "        self.squeeze = nn.Linear(in_channels, squeeze)\n",
    "        self.expand = nn.Linear(squeeze, in_channels)\n",
    "        self.activation = nn.ReLU()\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self._attention(x)\n",
    "\n",
    "    def _attention(self, x):\n",
    "        out = torch.mean(x, [2, 3])\n",
    "        out = self.squeeze(out)\n",
    "        out = self.activation(out)\n",
    "        out = self.expand(out)\n",
    "        out = self.sigmoid(out)\n",
    "        out = out.unsqueeze(2).unsqueeze(3)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Bottleneck(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_channels,\n",
    "        planes,\n",
    "        expansion,\n",
    "        stride=1,\n",
    "        cardinality=1,\n",
    "        se=False,\n",
    "        se_squeeze=16,\n",
    "        downsample=None,\n",
    "        fused_se=True,\n",
    "    ):\n",
    "        super(Bottleneck, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(in_channels, planes, 1, 1, int((1 - 1) / 2))\n",
    "        self.bn1 = nn.BatchNorm2d(planes)\n",
    "        self.conv2 = nn.Conv2d(planes, planes, 3, stride, int((3 - 1) / 2), groups=cardinality)\n",
    "        self.bn2 = nn.BatchNorm2d(planes)\n",
    "        self.conv3 = nn.Conv2d(planes, planes * expansion, 1, 1, int((1 - 1) / 2))\n",
    "        self.bn3 = nn.BatchNorm2d(planes * expansion)\n",
    "        self.relu = nn.ReLU()\n",
    "        \n",
    "        self.downsample = downsample\n",
    "        self.stride = stride\n",
    "\n",
    "        self.fused_se = fused_se\n",
    "        \n",
    "        if se:\n",
    "            self.squeeze = (\n",
    "                SqueezeAndExcitation(\n",
    "                    planes * expansion, se_squeeze\n",
    "                )\n",
    "            )\n",
    "        else:\n",
    "            self.squeeze = None\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv3(out)\n",
    "        out = self.bn3(out)\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            residual = self.downsample(x)\n",
    "\n",
    "        if self.squeeze is None:\n",
    "            out += residual\n",
    "        else:\n",
    "            if self.fused_se:\n",
    "                out = torch.addcmul(residual, out, self.squeeze(out), value=1)\n",
    "            else:\n",
    "                out = residual + out * self.squeeze(out)\n",
    "\n",
    "        out = self.relu(out)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResNet(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        image_channel: int = 3,\n",
    "        stem_width: int = 64,\n",
    "        layers=[3, 4, 23, 3],\n",
    "        widths=[128, 256, 512, 1024],\n",
    "        num_classes: int = 1000,\n",
    "        cardinality: int = 1,\n",
    "        expansion: int = 1,\n",
    "        last_bn_0_init: bool = False,\n",
    "        fused_se: bool = True,\n",
    "        se=False,\n",
    "        se_squeeze=16,\n",
    "    ):\n",
    "        super(ResNet, self).__init__()\n",
    "        self.last_bn_0_init = last_bn_0_init\n",
    "        self.conv1 = nn.Conv2d(image_channel, stem_width, 7, 2, int((7 - 1) / 2))\n",
    "        self.bn1 = nn.BatchNorm2d(stem_width)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "\n",
    "        self.se = se \n",
    "        self.se_squeeze = se_squeeze\n",
    "\n",
    "        in_channels = stem_width\n",
    "        assert len(widths) == len(layers)\n",
    "\n",
    "        self.layers = []\n",
    "        for i, (w, l) in enumerate(zip(widths, layers)):\n",
    "            layer, in_channels = self._make_layer(\n",
    "                expansion,\n",
    "                in_channels,\n",
    "                w,\n",
    "                l,\n",
    "                cardinality=cardinality,\n",
    "                stride=1 if i == 0 else 2,\n",
    "                fused_se=fused_se,\n",
    "                se=self.se,\n",
    "                se_squeeze=self.se_squeeze\n",
    "            )\n",
    "            self.layers.append(layer)\n",
    "\n",
    "        self.layers = nn.Sequential(*self.layers)\n",
    "        self.avgpool = nn.AdaptiveAvgPool2d(1)\n",
    "        self.fc = nn.Linear(widths[-1] * expansion, num_classes)\n",
    "\n",
    "    def stem(self, x):\n",
    "        x = self.conv1(x)\n",
    "        if self.bn1 is not None:\n",
    "            x = self.bn1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.maxpool(x)\n",
    "        return x\n",
    "\n",
    "    def classifier(self, x):\n",
    "        x = self.avgpool(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.fc(x)\n",
    "        return x\n",
    "\n",
    "    def _make_layer(\n",
    "            self,\n",
    "            expansion,\n",
    "            in_channels,\n",
    "            planes,\n",
    "            blocks,\n",
    "            se,\n",
    "            se_squeeze,\n",
    "            stride=1,\n",
    "            cardinality=1,\n",
    "            fused_se=True,\n",
    "        ):\n",
    "        downsample = None\n",
    "        if stride != 1 or in_channels != planes * expansion:\n",
    "            dconv = nn.Conv2d(in_channels, planes * expansion, 1, stride=stride, padding=0)\n",
    "            dbn = nn.BatchNorm2d(planes * expansion)\n",
    "            if dbn is not None:\n",
    "                downsample = nn.Sequential(dconv, dbn)\n",
    "            else:\n",
    "                downsample = dconv\n",
    "\n",
    "        layers = []\n",
    "        for i in range(blocks):\n",
    "            layers.append(\n",
    "                Bottleneck(\n",
    "                    in_channels=in_channels,\n",
    "                    planes=planes,\n",
    "                    expansion=expansion,\n",
    "                    stride=stride if i == 0 else 1,\n",
    "                    cardinality=cardinality,\n",
    "                    se=se,\n",
    "                    se_squeeze=se_squeeze,\n",
    "                    downsample=downsample if i == 0 else None,\n",
    "                    fused_se=fused_se,\n",
    "                )\n",
    "            )\n",
    "            in_channels = planes * expansion\n",
    "        return nn.Sequential(*layers), in_channels\n",
    "    \n",
    "    # todo add kaiming init\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.stem(x)\n",
    "        x = self.layers(x)\n",
    "        x = self.classifier(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet50 = ResNet(\n",
    "            image_channel=1,\n",
    "            stem_width=64,\n",
    "            layers=[3, 4, 6, 3],\n",
    "            widths=[64, 128, 256, 512],\n",
    "            expansion=4,\n",
    "            fused_se=True,\n",
    "            se=True\n",
    "        )\n",
    "\n",
    "resnet101 = ResNet(\n",
    "            image_channel=1,\n",
    "            stem_width=64,\n",
    "            layers=[3, 4, 23, 3],\n",
    "            widths=[128, 256, 512, 1024],\n",
    "            expansion=2,\n",
    "            cardinality=32,\n",
    "            fused_se=False\n",
    "        )\n",
    "\n",
    "resnet101se = ResNet(\n",
    "            image_channel=1,\n",
    "            stem_width=64,\n",
    "            layers=[3, 4, 23, 3],\n",
    "            widths=[128, 256, 512, 1024],\n",
    "            expansion=2,\n",
    "            cardinality=32,\n",
    "            fused_se=False\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inps = torch.randn((2, 1, 256, 256))\n",
    "\n",
    "out = resnet50(inps)\n",
    "\n",
    "out.size()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
